import os
import time

import numpy as np
from PIL import Image
from utils.plots import Annotator, colors

import yolov5


def tran(results):
    pred = results.pred
    ims = results.ims
    names = results.names

    labels = True

    s, crops = '', []
    for i, (im, pred) in enumerate(zip(ims, pred)):
        s += f'\nimage {i + 1}/{len(pred)}: {im.shape[0]}x{im.shape[1]} '  # string
        if pred.shape[0]:
            for c in pred[:, -1].unique():
                n = (pred[:, -1] == c).sum()  # detections per class
                s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
            annotator = Annotator(im, example=str(names))
            for *box, conf, cls in reversed(pred):  # xyxy, confidence, class
                label = f'{names[int(cls)]} {conf:.2f}'
                annotator.box_label(box, label if labels else '', color=colors(cls))
            im = annotator.im
        else:
            s += '(no detections)'

        img = np.array(im) if isinstance(im, Image.Image) else im
        return img


def getFiles(path):
    files = []
    for root, dirs, file in os.walk(path):
        for f in file:
            files.append(os.path.join(root, f))
    return files


if __name__ == '__main__':
    # model = yolov5.load('yolov5s.pt', device='0')
    model = yolov5.load('yolov5n.pt', device='0')
    files = getFiles('imgs')
    while True:
        for i, id in enumerate(files):
            time_begin = time.time()
            results = model(id)
            img = tran(results)
            print(i, time.time() - time_begin)
